#include "watermarking.h"

float alpha(int u) {
	return u? 0.5 : (float)sqrt(2)/4;
}

string get_watermarking_key() {
	printf("[%s] Insert private key for watermarking: ", timestamp());
	return getlongline(stdin);
}

dct_coeff_t dct(int k1, int k2, bmpfile_t* bmp) {
    int i, j; //i cycles ordinates, j cycles abscissæ
    float cos_1, cos_2, c=0;
    int min_i = floor(k1/BLOCK_SIZE)*BLOCK_SIZE;
    int min_j = floor(k2/BLOCK_SIZE)*BLOCK_SIZE;
    for(i=min_i; i<min_i+BLOCK_SIZE; i++) {
        cos_1 = cos(PI_BS*((i%BLOCK_SIZE)+0.5)*(k1%BLOCK_SIZE));
        for(j=min_j; j<min_j+BLOCK_SIZE; j++) {
            cos_2 = cos(PI_BS*((j%BLOCK_SIZE)+0.5)*(k2%BLOCK_SIZE));
            c += alpha(k1%BLOCK_SIZE) * alpha(k2%BLOCK_SIZE) * (rgb2mono(bmp->pixels[i][j])-128) * cos_1 * cos_2;
        }
    }
    return (dct_coeff_t)round(c/quant_table[k1%BLOCK_SIZE][k2%BLOCK_SIZE]);
}

rgb_pixel_t idct(int k1, int k2, int w, int h, dct_coeff_t coeff[w][h]) {
	int i, j; //i cycles ordinates, j cycles abscissæ
	float cos_1, cos_2, c=0;
	int min_i = floor(k1/BLOCK_SIZE)*BLOCK_SIZE;
    int min_j = floor(k2/BLOCK_SIZE)*BLOCK_SIZE;
    for(i=min_i; i<min_i+BLOCK_SIZE; i++) {
        cos_1 = cos(PI_BS*((k1%BLOCK_SIZE)+0.5)*(i%BLOCK_SIZE));
        for(j=min_j; j<min_j+BLOCK_SIZE; j++) {
            cos_2 = cos(PI_BS*((k2%BLOCK_SIZE)+0.5)*(j%BLOCK_SIZE));
            c += alpha(i%BLOCK_SIZE) * alpha(j%BLOCK_SIZE) * coeff[i][j] * quant_table[i%BLOCK_SIZE][j%BLOCK_SIZE] * cos_1 * cos_2;
        }
    }
    int ret = round(c) + 128;
    if(ret>255) ret = 255;
    if(ret<0) ret = 0;
    return (rgb_pixel_t){ret, ret, ret};
}

image embed(image source, string key) {
    printf("[%s] Watermark embedding started.\n", timestamp());

    int i, j, l;
    image img = clone(source, E_SUFFIX);
    int w = img->width;
    int h = img->height;
    dct_coeff_t coeff[w][h];
    int bytes = strlen(key) + 1;
    printf("[%s] Bytes to embed (including line terminator): %d.\n", timestamp(), bytes);
    printf("[%s] String to embed: \"%s\".\n", timestamp(), key);
    
    //1. fill DCT matrix
    printf("[%s] Calculating DCT coefficients...", timestamp());
    fflush(stdout);
    for(i=0; i<h; i++) {
        for(j=0; j<w; j++) {
			coeff[i][j] = dct(i, j, img->bitmap);
        }
    }
	printf("done.\n", timestamp());
	
    //2. apply ZRE using row-by-row scan order
    printf("[%s] ZRE started. Scan order used: row-by-row (default).\n", timestamp());
    int k = 0, t = 0, b = 8, bitmask, x, y;
    dct_coeff_t* triplet[3] = {NULL, NULL, NULL};
    for(i=0; i<h; i+=BLOCK_SIZE) {
        for(j=0; j<w; j+=BLOCK_SIZE) {
			for(x=i; x<i+BLOCK_SIZE; x++) {
				for(y=j; y<j+BLOCK_SIZE && k < bytes; y++) {
					if(x==i && y==j) continue; //skipping DC component
					triplet[t++] = &coeff[x][y];
					if(t==3) {
						t = 0;
						if(*triplet[0] && !*triplet[2]) {
							//confusion avoidance
							if(*triplet[1] > 0) {
								(*triplet[1])++;
							} else if(*triplet[1] < 0) {
								(*triplet[1])--;
							}
							//embedding
							if(!*triplet[1]) { //valid triplet found
								b--;
								bitmask = pow(2, b);
								*triplet[1] = (((unsigned int)key[k] & bitmask) == bitmask)? 1 : -1;
								if(!b) {
									b=8;
									k++;
								}
							}
						}
					}
				}
			}
        }
    }
    printf("[%s] Bytes embedded: %d.\n", timestamp(), k);
    
	//3. apply IDCT to watermarked DCT coefficients
	for(i=0; i<h; i++) {
        for(j=0; j<w; j++) {
			img->bitmap->pixels[i][j] = idct(i, j, w, h, coeff);
        }
    }
    
    printf("[%s] Watermarking complete.\n", timestamp());
    return img;
}

result extract(image marked) {
    printf("[%s] Watermark extraction started.\n", timestamp());
    
    int i, j, l = 0;
    image img = clone(marked, X_SUFFIX);
    int w = img->width;
    int h = img->height;
    dct_coeff_t coeff[w][h];
    result r;
	r.key = malloc((int)(w*h/3));

    //1. fill DCT matrix
    printf("[%s] Calculating DCT coefficients...", timestamp());
	fflush(stdout);
    for(i=0; i<h; i++) {
        for(j=0; j<w; j++) {
			coeff[i][j] = dct(i, j, img->bitmap);
        }
    }
	printf("done.\n", timestamp());
    
    //2. apply ZRX using row-by-row scan order
    printf("[%s] ZRX started. Scan order used: row-by-row (default).\n", timestamp());
    int k = 0, t = 0, string_found = 0, x, y;
    unsigned int c = 0;
    dct_coeff_t* triplet[3] = {NULL, NULL, NULL};
    for(i=0; i<h; i+=BLOCK_SIZE) {
        for(j=0; j<w; j+=BLOCK_SIZE) {
			for(x=i; x<i+BLOCK_SIZE; x++) {
				for(y=j; y<j+BLOCK_SIZE && !string_found; y++) {
					if(x==i && y==j) continue; //skipping DC component
					triplet[t++] = &coeff[x][y];
					if(t==3) {
						t = 0;
						if(*triplet[0] && !*triplet[2]) {
							//extracting
							if(*triplet[1]==1 || *triplet[1]==-1) {
								//watermark bit found
								if(*triplet[1]==1) {
									c += pow(2, (7 - k%8));
								}
								if(k%8==7) {
									r.key[l++] = (char)c;
									string_found = !c;
									c = 0;
								}
								k++;
								*triplet[1] = 0;
							} else if(*triplet[1] > 0) { //confusion avoidance
								(*triplet[1])--;
							} else if(*triplet[1] < 0) {
								(*triplet[1])++;
							}
						}
					}
				}
			}
        }
	}
    printf("[%s] Bytes extracted: %d.\n", timestamp(), strlen(r.key));
    
	//3. apply IDCT to unwatermarked DCT coefficients
	for(i=0; i<h; i++) {
        for(j=0; j<w; j++) {
			img->bitmap->pixels[i][j] = idct(i, j, w, h, coeff);
        }
    }
    
    printf("[%s] Watermark extraction complete.\n", timestamp());
    r.image = img;
    return r;
}
